use crate::tools::tool::random_from_stdnorm;
use nalgebra::linalg::Cholesky;
use nalgebra::{Matrix2, OMatrix, Vector2, U1, U2, U3};
use statrs::distribution::{Continuous, MultivariateNormal};
///Essentially what we are doing here is just simulating a standard geometric Brownian motion with non-constant volatility, where the change in S has relationship ρ with the change in volatility.
pub struct Heston {
    /// 初始资产价格
    pub S0: f64,
    pub T: f64,
    ///  long term mean  for dt
    pub mu: f64,
    /// 波动率调整速率
    kappa: f64,
    /// 长期平均波动率
    theta: f64,
    /// 初始波动率 方差形式
    sigma0: f64,
    /// 波动率和s的协方差
    rho: f64,
    volvol: f64,
    steps: usize,
}
/// formula : https://www.codearmo.com/python-tutorial/heston-model-simulation-python
#[doc = docify::embed!("src/stochastics/heston.rs", test_heston)]
impl Stochastic for Heston {
    fn gen_path(&self) -> Vec<f64> {
        let dt = self.T / self.steps as f64;
        let mut res = Vec::new();
        let mut v_t = self.sigma0;
        let mut s_t = self.S0;
        for step in 0..self.steps {
            /// wt是一个关于s和sigma的多元正太分布
            let (r1, r2) = gen_mvn2(self.rho);
            let wt1 = r1 * dt.sqrt();
            /// 价格方程
            let up = (self.mu - 0.5 * v_t) * dt + v_t.sqrt() * wt1;
            s_t = s_t * up.exp();
            /// 波动率方程
            let wt2 = r2 * dt.sqrt();
            v_t = v_t + self.kappa * (self.theta - v_t) * dt + self.volvol * v_t.sqrt() * wt2;
            v_t = v_t.abs(); // 应该是方差

            res.push(s_t);
        }
        res
    }
}

fn gen_mvn2(rho: f64) -> (f64, f64) {
    let mu = vec![0., 0.];
    let cov = vec![1., rho, rho, 1.];
    let mvn = MultivariateNormal::new(mu, cov).unwrap();
    let mut rng = rand::thread_rng();
    let res = mvn.sample(&mut rng);
    let res2 = (res[(0, 0)], res[(1, 0)]);
    res2
}
#[test]
#[docify::export]
fn test_heston() {
    let kappa = 4.;
    let theta = 0.02;
    let v_0 = 0.02;
    let volvol = 0.9;
    let r = 0.02;
    let S = 100.;
    let paths = 1;
    let steps = 20;
    let T = 1.;
    let rho = 0.5;
    let h = Heston {
        S0: S,
        T: T,
        mu: r,
        kappa: kappa,
        theta: theta,
        sigma0: v_0,
        rho: rho,
        volvol: volvol,
        steps: steps,
    };

    for i in 0..paths {
        let path = &h.gen_path();
        // dbg!(path);
    }
}
